#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
湖州地区真实数据采集器
使用高德地图API获取德清、长兴、安吉三个地方的街道、社区和小区真实数据
"""

import json
import os
import time
import sqlite3
import requests
from typing import List, Dict, Optional
import logging
import pandas as pd # Added missing import for pandas

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class HuzhouRealDataCollector:
    """湖州地区真实数据采集器"""
    
    def __init__(self):
        """初始化数据采集器"""
        self.api_key = os.getenv('AMAP_WEB_KEY', '85477d3ecd81e43af7f4e931912c3f43')
        self.base_url = 'https://restapi.amap.com/v3'
        self.db_path = "huzhou_real_data.db"
        self.output_dir = "data"
        
        # 湖州地区配置
        self.districts = {
            '330521': '德清县',
            '330522': '长兴县', 
            '330523': '安吉县'
        }
        
        # 创建输出目录
        os.makedirs(self.output_dir, exist_ok=True)
        
        # 初始化数据库
        self.init_database()
    
    def init_database(self):
        """初始化数据库"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 创建街道表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS streets (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                adcode TEXT UNIQUE,
                name TEXT,
                center TEXT,
                polyline TEXT,
                district_name TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建社区表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS communities (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                adcode TEXT UNIQUE,
                name TEXT,
                center TEXT,
                polyline TEXT,
                street_name TEXT,
                district_name TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建小区表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS residential_areas (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT,
                address TEXT,
                location TEXT,
                community_name TEXT,
                street_name TEXT,
                district_name TEXT,
                poi_id TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        conn.commit()
        conn.close()
        logger.info("数据库初始化完成")
    
    def fetch_district_data(self, adcode: str, district_name: str) -> Dict:
        """获取指定地区的街道数据"""
        url = f"{self.base_url}/config/district"
        params = {
            'keywords': adcode,
            'subdistrict': 1,  # 获取下一级行政区划（街道）
            'extensions': 'all',
            'key': self.api_key
        }
        
        try:
            logger.info(f"正在获取 {district_name} ({adcode}) 的街道数据...")
            resp = requests.get(url, params=params, timeout=10)
            resp.raise_for_status()
            data = resp.json()
            
            if data.get('status') != '1':
                raise RuntimeError(f"API错误 {adcode}: {data.get('info')}")
            
            district_list = data.get('districts', [])
            if not district_list:
                return {'name': district_name, 'streets': []}
            
            streets = []
            for street in district_list[0].get('districts', []):
                streets.append({
                    'name': street.get('name'),
                    'adcode': street.get('adcode'),
                    'center': street.get('center'),
                    'polyline': street.get('polyline')
                })
            
            return {
                'name': district_name,
                'streets': streets
            }
            
        except Exception as e:
            logger.error(f"获取 {district_name} 数据失败: {e}")
            return {'name': district_name, 'streets': []}
    
    def fetch_community_data(self, street_adcode: str, street_name: str, district_name: str) -> List[Dict]:
        """获取指定街道的社区数据"""
        url = f"{self.base_url}/config/district"
        params = {
            'keywords': street_adcode,
            'subdistrict': 1,  # 获取下一级行政区划（社区）
            'extensions': 'all',
            'key': self.api_key
        }
        
        try:
            logger.info(f"正在获取 {street_name} 的社区数据...")
            resp = requests.get(url, params=params, timeout=10)
            resp.raise_for_status()
            data = resp.json()
            
            if data.get('status') != '1':
                logger.warning(f"获取社区数据失败 {street_name}: {data.get('info')}")
                return []
            
            district_list = data.get('districts', [])
            if not district_list:
                return []
            
            communities = []
            for community in district_list[0].get('districts', []):
                communities.append({
                    'name': community.get('name'),
                    'adcode': community.get('adcode'),
                    'center': community.get('center'),
                    'polyline': community.get('polyline'),
                    'street_name': street_name,
                    'district_name': district_name
                })
            
            return communities
            
        except Exception as e:
            logger.error(f"获取 {street_name} 社区数据失败: {e}")
            return []
    
    def fetch_residential_areas(self, community_name: str, street_name: str, district_name: str) -> List[Dict]:
        """获取指定社区的小区数据（通过POI搜索）"""
        url = f"{self.base_url}/place/text"
        params = {
            'keywords': f"{community_name}小区",
            'city': district_name,
            'types': '商务住宅',
            'extensions': 'all',
            'key': self.api_key,
            'offset': 50,  # 每页50条
            'page': 1
        }
        
        try:
            logger.info(f"正在获取 {community_name} 的小区数据...")
            resp = requests.get(url, params=params, timeout=10)
            resp.raise_for_status()
            data = resp.json()
            
            if data.get('status') != '1':
                logger.warning(f"获取小区数据失败 {community_name}: {data.get('info')}")
                return []
            
            pois = data.get('pois', [])
            residential_areas = []
            
            for poi in pois:
                # 过滤出真正的小区
                if any(keyword in poi.get('name', '') for keyword in ['小区', '花园', '公寓', '家园', '广场', '大厦']):
                    residential_areas.append({
                        'name': poi.get('name'),
                        'address': poi.get('address'),
                        'location': poi.get('location'),
                        'community_name': community_name,
                        'street_name': street_name,
                        'district_name': district_name,
                        'poi_id': poi.get('id')
                    })
            
            return residential_areas
            
        except Exception as e:
            logger.error(f"获取 {community_name} 小区数据失败: {e}")
            return []
    
    def save_streets(self, district_name: str, streets: List[Dict]):
        """保存街道信息到数据库"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        for street in streets:
            try:
                cursor.execute('''
                    INSERT OR REPLACE INTO streets (adcode, name, center, polyline, district_name)
                    VALUES (?, ?, ?, ?, ?)
                ''', (
                    street.get("adcode"),
                    street.get("name"),
                    street.get("center"),
                    street.get("polyline"),
                    district_name
                ))
            except Exception as e:
                logger.error(f"保存街道信息失败: {e}")
        
        conn.commit()
        conn.close()
        logger.info(f"保存了 {len(streets)} 条街道信息")
    
    def save_communities(self, communities: List[Dict]):
        """保存社区信息到数据库"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        for community in communities:
            try:
                cursor.execute('''
                    INSERT OR REPLACE INTO communities (adcode, name, center, polyline, street_name, district_name)
                    VALUES (?, ?, ?, ?, ?, ?)
                ''', (
                    community.get("adcode"),
                    community.get("name"),
                    community.get("center"),
                    community.get("polyline"),
                    community.get("street_name"),
                    community.get("district_name")
                ))
            except Exception as e:
                logger.error(f"保存社区信息失败: {e}")
        
        conn.commit()
        conn.close()
        logger.info(f"保存了 {len(communities)} 条社区信息")
    
    def save_residential_areas(self, residential_areas: List[Dict]):
        """保存小区信息到数据库"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        for area in residential_areas:
            try:
                cursor.execute('''
                    INSERT OR REPLACE INTO residential_areas (name, address, location, community_name, street_name, district_name, poi_id)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                ''', (
                    area.get("name"),
                    area.get("address"),
                    area.get("location"),
                    area.get("community_name"),
                    area.get("street_name"),
                    area.get("district_name"),
                    area.get("poi_id")
                ))
            except Exception as e:
                logger.error(f"保存小区信息失败: {e}")
        
        conn.commit()
        conn.close()
        logger.info(f"保存了 {len(residential_areas)} 条小区信息")
    
    def export_to_json(self):
        """导出数据到JSON文件"""
        conn = sqlite3.connect(self.db_path)
        
        # 获取街道数据
        streets_df = pd.read_sql_query("SELECT * FROM streets", conn)
        communities_df = pd.read_sql_query("SELECT * FROM communities", conn)
        areas_df = pd.read_sql_query("SELECT * FROM residential_areas", conn)
        
        conn.close()
        
        data = {
            'streets': streets_df.to_dict('records'),
            'communities': communities_df.to_dict('records'),
            'residential_areas': areas_df.to_dict('records')
        }
        
        output_file = os.path.join(self.output_dir, 'huzhou_real_data.json')
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        logger.info(f"数据已导出到 {output_file}")
    
    def collect_all_data(self):
        """采集所有数据"""
        logger.info("开始采集湖州地区真实数据...")
        
        all_data = {}
        
        for adcode, district_name in self.districts.items():
            print(f"\n{'='*50}")
            print(f"正在处理 {district_name} ({adcode})")
            print(f"{'='*50}")
            
            # 获取街道数据
            district_data = self.fetch_district_data(adcode, district_name)
            all_data[adcode] = district_data
            
            # 保存街道数据
            self.save_streets(district_name, district_data['streets'])
            
            # 获取每个街道的社区数据
            for street in district_data['streets']:
                street_name = street['name']
                street_adcode = street['adcode']
                
                print(f"\n  📂 处理街道: {street_name}")
                
                # 获取社区数据
                communities = self.fetch_community_data(street_adcode, street_name, district_name)
                self.save_communities(communities)
                
                # 获取每个社区的小区数据
                for community in communities:
                    community_name = community['name']
                    print(f"    📄 处理社区: {community_name}")
                    
                    # 获取小区数据
                    residential_areas = self.fetch_residential_areas(community_name, street_name, district_name)
                    self.save_residential_areas(residential_areas)
                    
                    # 控制请求频率
                    time.sleep(0.2)
                
                # 街道间请求间隔
                time.sleep(0.5)
            
            # 地区间请求间隔
            time.sleep(1)
        
        # 导出数据
        self.export_to_json()
        
        logger.info("数据采集完成！")
        return all_data

def main():
    """主函数"""
    collector = HuzhouRealDataCollector()
    
    # 检查API密钥
    if not collector.api_key:
        print("❌ 未配置高德地图API密钥")
        print("请设置环境变量: export AMAP_WEB_KEY='your_api_key'")
        return
    
    print("🏘️ 湖州地区真实数据采集器")
    print("="*50)
    print(f"API密钥: {'✅ 已配置' if collector.api_key else '❌ 未配置'}")
    print(f"目标地区: {', '.join(collector.districts.values())}")
    print("="*50)
    
    # 开始采集数据
    try:
        all_data = collector.collect_all_data()
        print("\n✅ 数据采集完成！")
        print(f"📁 数据文件: {collector.db_path}")
        print(f"📄 JSON文件: {os.path.join(collector.output_dir, 'huzhou_real_data.json')}")
        
    except KeyboardInterrupt:
        print("\n⚠️ 用户中断采集")
    except Exception as e:
        print(f"\n❌ 采集失败: {e}")

if __name__ == "__main__":
    main() 